import os
import numpy as np
import re

def merge_flow_fields(flow_dir, output_file, tile_size=256):
    # 获取所有光流场文件
    flow_files = [f for f in os.listdir(flow_dir) if f.endswith('.npy')]

    # 解析所有光流场文件的坐标
    coordinates = []
    for flow_file in flow_files:
        match = re.search(r'tile_\((\d+)_(\d+)\)\.npy', flow_file)
        if not match:
            raise ValueError(f"文件名格式不正确: {flow_file}")
        row, col = map(int, match.groups())
        coordinates.append((row, col))

    # 确定大光流场的尺寸
    max_row = max(row for row, col in coordinates)
    max_col = max(col for row, col in coordinates)
    large_flow = np.zeros(((max_row + 1) * tile_size, (max_col + 1) * tile_size, 2), dtype=np.float32)

    # 填充大光流场
    for flow_file in flow_files:
        flow_path = os.path.join(flow_dir, flow_file)
        flow = np.load(flow_path)

        match = re.search(r'tile_\((\d+)_(\d+)\)\.npy', flow_file)
        row, col = map(int, match.groups())

        start_row, start_col = row * tile_size, col * tile_size
        large_flow[start_row:start_row + tile_size, start_col:start_col + tile_size] = flow

    # 创建输出目录（如果不存在）
    output_dir = os.path.dirname(output_file)
    os.makedirs(output_dir, exist_ok=True)

    # 保存大光流场
    np.save(output_file, large_flow)

# 示例用法
flow_dir = r'E:\wafer52\256nm_flow_split_128'  # 替换为实际路径
output_file = r'E:\wafer52\256nm_flow_combine\large_flow_repair3.npy'  # 替换为实际路径

merge_flow_fields(flow_dir, output_file,tile_size=128)

print(f"Processed large flow field saved to {output_file}")
